import asyncio
from dataclasses import dataclass, asdict
from typing import Literal, List, Dict, Optional
from agents import Agent, Runner, OpenAIProvider, RunResult, RunConfig, set_tracing_export_api_key, set_tracing_disabled
from openai import AsyncOpenAI
import os
import json

# 实验提供的 API 配置
BASE_URL = "https://api.chatfire.cn/v1"
API_KEY = "sk-UcrTOYgPf03qqxdGspmJtaUmP7OTkkzuxsRHWPpMSxKXK8IN"
os.environ["OPENAI_API_KEY"] = API_KEY

# 配置Tracing客户端使用相同的API密钥
set_tracing_export_api_key(API_KEY)
set_tracing_disabled(True)


# 定义需求模型 DSL 结构
@dataclass
class UseCase:
    name: str
    include: List[str]
    extend: List[str]

@dataclass
class SystemSequenceDiagram:
    name: str
    steps: List[str]

@dataclass
class ConceptClass:
    name: str
    attributes: List[str]
    methods: List[str]

@dataclass
class OCLContract:
    context: str
    constraints: List[str]

@dataclass
class RequirementModel:
    use_case_diagram: List[UseCase]
    sequence_diagrams: List[SystemSequenceDiagram]
    concept_classes: List[ConceptClass]
    ocl_contracts: List[OCLContract]


# 定义评审反馈格式
@dataclass
class EvaluationFeedback:
    score: Literal["pass", "needs_improvement", "fail"]
    feedback: str


# 定义建模专家 Agent
modeling_expert = Agent(
    name="ModelingExpert",
    model="gpt-4o",
    instructions=(
        "您是一位需求建模专家，擅长根据用户需求创建完整的需求模型。"
        "请根据用户描述的需求，生成包含以下内容的需求模型：\n"
        "1. 用例图（包含用例、参与者及其关系）\n"
        "2. 系统顺序图（关键交互流程）\n"
        "3. 概念类图（主要领域对象及其关系）\n"
        "4. OCL合约（关键业务规则约束）\n\n"
        "输出必须严格使用指定的 RequirementModel DSL 格式。"
        "如果收到评审反馈，请根据反馈改进模型。"
    ),
    output_type=RequirementModel  # 指定输出格式
)

# 定义评审专家 Agent
model_reviewer = Agent(
    name="ModelReviewer",
    model="gpt-4o",
    instructions=(
        "您是一位需求模型评审专家，负责评估建模专家创建的需求模型质量。\n"
        "评审标准：\n"
        "1. 完整性：是否覆盖所有用户需求\n"
        "2. 一致性：不同模型间是否存在冲突\n"
        "3. 准确性：是否符合领域知识\n"
        "4. 可追溯性：需求与模型元素是否可对应\n\n"
        "对于不通过的模型，必须提供具体的改进建议，改进建议请用中文回答。"
        "输出必须使用 EvaluationFeedback 格式。"
    ),
    output_type=EvaluationFeedback  # 指定输出格式
)

# 配置 OpenAI 服务提供者
provider = OpenAIProvider(
    openai_client=AsyncOpenAI(base_url=BASE_URL, api_key=API_KEY),
    use_responses=False
)


# 定义需求建模工作流
async def requirement_modeling_workflow(user_requirement: str) -> RequirementModel:
    input_items = [{"content": user_requirement, "role": "user"}]
    current_model = None
    iteration = 0
    max_iterations = 2  # 最大迭代次数

    # 使用 trace 记录工作流执行
    # with trace("RequirementModelingWorkflow"):
    while iteration < max_iterations:
        # 步骤1: 生成需求模型
        modeling_result = await Runner.run(
            modeling_expert,
            input_items,
            run_config=RunConfig(model_provider=provider)
        )

        current_model = modeling_result.final_output
        print(f"\n=== 第 {iteration + 1} 次迭代生成的需求模型 ===")
        print(current_model)

        # 将模型添加到对话上下文
        input_items.extend(modeling_result.to_input_list())

        # 步骤2: 评审模型
        review_result = await Runner.run(
            model_reviewer,
            input_items,
            run_config=RunConfig(model_provider=provider)
        )

        feedback = review_result.final_output
        print(f"\n评审结果: {feedback.score}")
        print(f"反馈意见: {feedback.feedback}")

        # 检查是否通过
        if feedback.score == "pass":
            print("\n模型评审通过!")
            return current_model

        # 未通过则添加反馈继续迭代
        input_items.append({
            "content": f"模型评审反馈: {feedback.feedback}",
            "role": "user"
        })
        iteration += 1

    print("\n达到最大迭代次数，返回最终模型")
    return current_model


# 主函数
async def main():
    # 用户需求描述（示例）
    user_requirement = """
    我需要一个脑电信号采集系统，功能包括：
    1. 医生可采集患者的脑电信号
    2. 医生可创建及管理脑电分析报告
    3. 患者可查看或下载脑电分析报告
    4. 管理员理可管患者信息，医生信息
    5. 患者需要先支付脑电检查费用
    """

    print("开始需求建模工作流...")
    final_model = await requirement_modeling_workflow(user_requirement)

    print("\n=== 最终需求模型 ===")
    print(final_model)
    # 保存为 JSON 文件
    with open("requirement_model.json", "w") as f:
        json.dump(asdict(final_model), f, indent=2)

    print("模型已保存到 requirement_model.json")


if __name__ == "__main__":
    asyncio.run(main())